package in.ac.iitb.cse.nlp.postagger;

import in.ac.iitb.cse.nlp.postagger.data.BiViterbiHMM;
import in.ac.iitb.cse.nlp.postagger.data.DebugConsole;
import in.ac.iitb.cse.nlp.postagger.data.EmissionMatrix;
import in.ac.iitb.cse.nlp.postagger.data.TransitionMatrix;

import java.io.BufferedReader;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

public class TableCreator {

	// For accuracy testing
	long noOfTotalPredictions;
	long noOfCorrectPredictions;

	static final String WORD_SEPARATOR = " ";
	static final String TAG_SEPARATOR = "_";
	// We will divide the input into sentences so that
	// we can process it faster
	static final String START_OF_SENTENCE = "^";
	static final String END_OF_SENTENCE = ".";
	static final String ALTERNATE_END_OF_SENTENCE = "?";
	static Map<String, HashMap<String, Integer>> transitionTable = new HashMap<String, HashMap<String, Integer>>();
	static Map<String, HashMap<String, Integer>> emissionTable = new HashMap<String, HashMap<String, Integer>>();

	// For new table creation
	static TransitionMatrix transitionMatrix = new TransitionMatrix();
	static EmissionMatrix emissionMatrix = new EmissionMatrix();
	static Set<String> tagSet = new HashSet<String>();

	public void createNew(File file) throws IOException {
		BufferedReader reader = new BufferedReader(new InputStreamReader(
				new FileInputStream(file)));
		String line = null;
		tagSet.add(START_OF_SENTENCE);
		while ((line = reader.readLine()) != null) {
			String[] split = line.split(WORD_SEPARATOR);
			String previousTag = START_OF_SENTENCE;
			for (String string : split) {
				String[] strings = string.split(TAG_SEPARATOR);
				String currentTag = strings[0];
				tagSet.add(currentTag);
				// For updating transition table
				if (previousTag != null) {
					transitionMatrix.update(previousTag, currentTag);
				}
				// For updating emission table
				emissionMatrix.update(currentTag, strings[1]);
				if (strings[1].equals(END_OF_SENTENCE)
						|| strings[1].equals(ALTERNATE_END_OF_SENTENCE)) {
					previousTag = START_OF_SENTENCE;
				} else {
					previousTag = currentTag;
				}
			}
		}
	}

	public void newDisplayTables() {
		System.out.println(transitionMatrix.toString());
		System.out.println(emissionMatrix.toString());
		System.out.println(tagSet);
		System.out.println(tagSet.size());
	}

	public void go() {
		File file = new File("corpus/A00.txt");
		try {
			createNew(file);
			newDisplayTables();
			String sentence = "^ I am happy . $";
			String[] split = sentence.split(" ");
			List<String> input = new ArrayList<String>();
			for (String string : split)
				input.add(string);
			BiViterbiHMM vhmm = new BiViterbiHMM();
			vhmm.init(new ArrayList<String>(tagSet), transitionMatrix,
					emissionMatrix);
			List<String> bestProbTagSeq = vhmm.bestProbTagSeq(input);
			System.out.println(bestProbTagSeq);
		} catch (FileNotFoundException e) {
			e.printStackTrace();
		} catch (IOException e) {
			e.printStackTrace();
		}
	}

	public void goFile(String filename) {
		if (filename.endsWith("txt") == false)
			return;
		System.out.println(filename + " is Trained");
		File file = new File(filename);
		try {
			createNew(file);
		} catch (FileNotFoundException e) {
			e.printStackTrace();
		} catch (IOException e) {
			e.printStackTrace();
		}
	}

	public void goDir(String cdir, String dir) {
		if (cdir.endsWith("/") == false && dir.endsWith("/") == false)
			cdir.concat("/");
		File Dir = new File(cdir + dir);
		if (Dir.isDirectory()) {
			for (File f : Dir.listFiles()) {
				// Dont assume that file trained good...
				// return the exit Status, check it
				if (f.isFile() && f.canRead()) {
					goFile(f.getAbsolutePath());
				}
			}

		} else if (Dir.canRead() && Dir.isFile()) {
			// Dont assume that file trained good...
			// return the exit Status, check it
			goFile(Dir.getAbsolutePath());
		}
	}

	public double fiveFold(String cdir, String dir) throws IOException {
		if (cdir.endsWith("/") == false && dir.endsWith("/") == false)
			cdir.concat("/");
		File Dir = new File(cdir + dir);
		if (Dir.isDirectory()) {
			// First do training with the 80% of the files
			File[] Files = Dir.listFiles();
			int trainingFiles = (int) (Files.length * 0.8);
			for (int i = 0; i < trainingFiles; i++) {
				File file = Files[i];
				if (file.isFile() && file.canRead()) {
					goFile(file.getAbsolutePath());
				}
			}
			// Now test with 20% of the files whether the tags
			// are correct
			for (int i = trainingFiles; i < Files.length; i++) {
				File file = Files[i];
				if (file.isFile() && file.canRead()) {
					test(file);
				}
				System.out
						.println(file.getAbsolutePath() + " used for testing");
			}
		}
		System.out.println(noOfCorrectPredictions);
		System.out.println(noOfTotalPredictions);
		return (((double) noOfCorrectPredictions) / (double) (noOfTotalPredictions));
	}

	private void test(File file) throws IOException {
		BufferedReader reader = new BufferedReader(new InputStreamReader(
				new FileInputStream(file)));
		String line = null;
		while ((line = reader.readLine()) != null) {
			String[] split = line
					.split("PUN" + TAG_SEPARATOR + END_OF_SENTENCE);
			for (String string : split) {
				String[] split2 = string.split("PUN" + TAG_SEPARATOR
						+ ALTERNATE_END_OF_SENTENCE);
				if (split2.length > 1) {
					predict(split2[0], "PUN" + TAG_SEPARATOR
							+ ALTERNATE_END_OF_SENTENCE);
					predict(split2[1], "PUN" + TAG_SEPARATOR + END_OF_SENTENCE);
				} else {
					predict(split2[0], "PUN" + TAG_SEPARATOR + END_OF_SENTENCE);
				}
			}
		}
	}

	private void predict(String string, String endOfSentence) {
		string = string.trim() + " " + endOfSentence;
		String[] split = string.trim().split(WORD_SEPARATOR);
		List<String> actualTagSequence = new ArrayList<String>();
		actualTagSequence.add(START_OF_SENTENCE);
		List<String> sentence = new ArrayList<String>();
		sentence.add("^");
		for (String splitString : split) {
			String[] split2 = splitString.split(TAG_SEPARATOR);
			actualTagSequence.add(split2[0]);
			sentence.add(split2[1]);
		}
		actualTagSequence.add("$");
		sentence.add(" $");
		BiViterbiHMM vhmm = new BiViterbiHMM();
		vhmm.init(new ArrayList<String>(tagSet), transitionMatrix,
				emissionMatrix);
		List<String> bestProbTagSeq = vhmm.bestProbTagSeq(sentence);
		if (bestProbTagSeq != null
				&& bestProbTagSeq.size() == actualTagSequence.size()) {
			for (int i = 1; i < actualTagSequence.size() - 1; i++) {
				noOfTotalPredictions++;
				if (actualTagSequence.get(i).equals(bestProbTagSeq.get(i))) {
					noOfCorrectPredictions++;
				}
			}
		}
		System.out.print("");
	}

	public List<String> tagSentence(String sentence) {
		String[] split = sentence.split(" ");
		List<String> input = new ArrayList<String>();
		for (String string : split)
			input.add(string);
		BiViterbiHMM vhmm = new BiViterbiHMM();
		vhmm.init(new ArrayList<String>(tagSet), transitionMatrix,
				emissionMatrix);
		List<String> bestProbTagSeq = vhmm.bestProbTagSeq(input);
		return bestProbTagSeq;
	}

	public static void main(String[] args) throws IOException {
		TableCreator tableCreator = new TableCreator();
		tableCreator.goDir("", "corpus");
		DebugConsole dbc = new DebugConsole();
		dbc.start(TableCreator.transitionMatrix, TableCreator.emissionMatrix,
				new ArrayList<String>(TableCreator.tagSet));
	}
}